#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <gmssl/sm2.h>
#include <gmssl/sm3.h>
#include <gmssl/sm4.h>
#include <gmssl/error.h>
#include <gmssl/rand.h>
#include <gmssl/hex.h>
#include <gmssl/sm2_z256.h>

#define MAX_PLAINTEXT_SIZE 1024
#define MAX_CIPHERTEXT_SIZE 1280

// 辅助函数：打印十六进制数据
void print_hex(const char *label, const uint8_t *data, size_t len)
{
    printf("%s: ", label);
    for (size_t i = 0; i < len; i++)
    {
        printf("%02X", data[i]);
        if ((i + 1) % 16 == 0 && i < len - 1)
        {
            printf("\n   ");
        }
        else if ((i + 1) % 4 == 0 && i < len - 1)
        {
            printf(" ");
        }
    }
    printf("\n");
}

int main()
{
    int ret = -1;

    printf("========== 数字信封协议实现 ==========\n\n");

    // 原始明文
    uint8_t plaintext[] = "这是一条需要安全传输的机密消息。使用数字信封保护！";
    size_t plaintext_len = strlen((const char *)plaintext);
    printf("原始明文: %s\n", plaintext);
    printf("明文长度: %zu 字节\n\n", plaintext_len);

    // Step 1: 生成SM2密钥对（发送方A和接收方B）
    printf("步骤 1: 生成SM2密钥对...\n");
    SM2_KEY sm2_key_a, sm2_key_b;
    uint8_t sm2_public_key_a_raw[65];
    uint8_t sm2_public_key_b_raw[65];
    size_t sm2_public_key_a_raw_len, sm2_public_key_b_raw_len;

    if (sm2_key_generate(&sm2_key_a) != 1)
    {
        fprintf(stderr, "Error: 发送方SM2密钥生成失败\n");
        return 1;
    }
    if (sm2_key_generate(&sm2_key_b) != 1)
    {
        fprintf(stderr, "Error: 接收方SM2密钥生成失败\n");
        return 1;
    }

    // 提取公钥用于后续验证
    sm2_z256_point_to_uncompressed_octets(&sm2_key_a.public_key, sm2_public_key_a_raw);
    sm2_z256_point_to_uncompressed_octets(&sm2_key_b.public_key, sm2_public_key_b_raw);

    sm2_public_key_a_raw_len = 65; // 未压缩的SM2公钥长度为65字节
    sm2_public_key_b_raw_len = 65;

    print_hex("发送方A的SM2公钥", sm2_public_key_a_raw, sm2_public_key_a_raw_len);
    print_hex("接收方B的SM2公钥", sm2_public_key_b_raw, sm2_public_key_b_raw_len);
    printf("\n");

    // Step 2: 随机生成SM4密钥
    printf("步骤 2: 随机生成SM4密钥...\n");
    uint8_t sm4_key[SM4_KEY_SIZE];
    SM4_KEY sm4_enc_key;

    rand_bytes(sm4_key, sizeof(sm4_key));
    print_hex("随机生成的SM4密钥", sm4_key, sizeof(sm4_key));

    // 使用SM4密钥设置加密密钥
    sm4_set_encrypt_key(&sm4_enc_key, sm4_key);
    printf("\n");

    // Step 3: 使用SM4密钥加密明文
    printf("步骤 3: 使用SM4密钥加密明文...\n");
    uint8_t sm4_ciphertext[MAX_CIPHERTEXT_SIZE];
    size_t sm4_ciphertext_len;

    // 计算填充后的长度
    sm4_ciphertext_len = (plaintext_len / SM4_BLOCK_SIZE + 1) * SM4_BLOCK_SIZE;

    // 简单实现: 使用ECB模式加密（实际应用中应考虑使用CBC、GCM等更安全的模式）
    for (size_t i = 0; i < plaintext_len / SM4_BLOCK_SIZE + 1; i++)
    {
        uint8_t block[SM4_BLOCK_SIZE] = {0}; // 初始化为0用于填充
        size_t copy_len = (i == plaintext_len / SM4_BLOCK_SIZE) ? plaintext_len % SM4_BLOCK_SIZE : SM4_BLOCK_SIZE;

        if (copy_len > 0)
        {
            memcpy(block, plaintext + i * SM4_BLOCK_SIZE, copy_len);
        }

        sm4_encrypt(&sm4_enc_key, block, sm4_ciphertext + i * SM4_BLOCK_SIZE);
    }

    print_hex("SM4加密密文", sm4_ciphertext, sm4_ciphertext_len);
    printf("\n");

    // Step 4: 对SM4密文计算SM3摘要
    printf("步骤 4: 对SM4密文计算SM3摘要...\n");
    uint8_t sm3_digest[SM3_DIGEST_SIZE];
    SM3_CTX sm3_ctx;

    sm3_init(&sm3_ctx);
    sm3_update(&sm3_ctx, sm4_ciphertext, sm4_ciphertext_len);
    sm3_finish(&sm3_ctx, sm3_digest);

    print_hex("SM4密文的SM3摘要", sm3_digest, sizeof(sm3_digest));
    printf("\n");

    // Step 5: 使用发送方的SM2私钥对摘要签名
    printf("步骤 5: 使用发送方的SM2私钥对摘要签名...\n");
    uint8_t sm2_signature[SM2_MAX_SIGNATURE_SIZE];
    size_t sm2_signature_len = 0;

    if (sm2_sign(&sm2_key_a, sm3_digest, sm2_signature, &sm2_signature_len) != 1)
    {
        fprintf(stderr, "Error: sm2_sign failed\n");
        return 1;
    }
    print_hex("SM2 Signature", sm2_signature, sm2_signature_len);
    printf("\n");

    // Step 6: 使用接收方的SM2公钥加密SM4密钥
    printf("步骤 6: 使用接收方的SM2公钥加密SM4密钥...\n");
    uint8_t encrypted_sm4_key[SM2_MAX_CIPHERTEXT_SIZE];
    size_t encrypted_sm4_key_len = 0;

    if (sm2_encrypt(&sm2_key_b, sm4_key, sizeof(sm4_key),
                    encrypted_sm4_key, &encrypted_sm4_key_len) != 1)
    {
        fprintf(stderr, "Error: SM2加密SM4密钥失败\n");
        return 1;
    }

    print_hex("SM2加密后的SM4密钥", encrypted_sm4_key, encrypted_sm4_key_len);
    printf("\n");

    // 至此，数字信封的组成部分已生成完毕
    // 数字信封包括：SM4密文 + SM2加密的SM4密钥 + SM2签名
    printf("步骤 7: 生成完整的数字信封...\n");
    printf("数字信封组成部分:\n");
    printf("- SM4加密的原始数据密文\n");
    printf("- SM2加密的SM4密钥\n");
    printf("- SM2签名(保证数据完整性和发送方身份)\n\n");

    // ========================= 接收方解封数字信封 =========================
    printf("========== 接收方解封数字信封 ==========\n\n");

    // Step 1: 使用接收方的SM2私钥解密SM4密钥
    printf("步骤 1: 使用接收方的SM2私钥解密SM4密钥...\n");
    uint8_t decrypted_sm4_key[SM4_KEY_SIZE];
    size_t decrypted_sm4_key_len = 0;

    if (sm2_decrypt(&sm2_key_b, encrypted_sm4_key, encrypted_sm4_key_len,
                    decrypted_sm4_key, &decrypted_sm4_key_len) != 1)
    {
        fprintf(stderr, "Error: SM2解密SM4密钥失败\n");
        return 1;
    }

    print_hex("解密后的SM4密钥", decrypted_sm4_key, decrypted_sm4_key_len);

    // 验证解密的SM4密钥是否正确
    if (decrypted_sm4_key_len != sizeof(sm4_key) ||
        memcmp(decrypted_sm4_key, sm4_key, sizeof(sm4_key)) != 0)
    {
        fprintf(stderr, "Error: 解密的SM4密钥与原始密钥不匹配\n");
        return 1;
    }
    printf("SM4密钥解密成功，与原始密钥匹配。\n\n");

    // Step 2: 对接收到的SM4密文计算SM3摘要
    printf("步骤 2: 对接收到的SM4密文计算SM3摘要...\n");
    uint8_t received_sm3_digest[SM3_DIGEST_SIZE];
    SM3_CTX received_sm3_ctx;

    sm3_init(&received_sm3_ctx);
    sm3_update(&received_sm3_ctx, sm4_ciphertext, sm4_ciphertext_len);
    sm3_finish(&received_sm3_ctx, received_sm3_digest);

    print_hex("接收到的SM4密文的SM3摘要", received_sm3_digest, sizeof(received_sm3_digest));
    printf("\n");

    // Step 3: 使用发送方的SM2公钥验证签名
    printf("步骤 3: 使用发送方的SM2公钥验证签名...\n");
    if (sm2_verify(&sm2_key_a, received_sm3_digest, sm2_signature, sm2_signature_len) != 1)
    {
        fprintf(stderr, "Error: sm2_verify failed\n");
        return 1;
    }
    printf("签名验证成功。数据完整性和发送方身份已确认。\n\n");

    // Step 4: 使用解密的SM4密钥解密密文
    printf("步骤 4: 使用解密的SM4密钥解密密文...\n");
    uint8_t decrypted_plaintext[MAX_PLAINTEXT_SIZE] = {0};
    SM4_KEY sm4_dec_key;

    sm4_set_decrypt_key(&sm4_dec_key, decrypted_sm4_key);

    for (size_t i = 0; i < sm4_ciphertext_len / SM4_BLOCK_SIZE; i++)
    {
        sm4_encrypt(&sm4_dec_key, sm4_ciphertext + i * SM4_BLOCK_SIZE, decrypted_plaintext + i * SM4_BLOCK_SIZE);
    }

    printf("解密后的明文: %s\n", decrypted_plaintext);

    // 验证解密后的明文是否与原始明文一致
    if (memcmp(decrypted_plaintext, plaintext, plaintext_len) == 0)
    {
        printf("解密成功！解密后的明文与原始明文一致。\n\n");
        ret = 0;
    }
    else
    {
        fprintf(stderr, "Error: 解密后的明文与原始明文不一致\n");
    }

    printf("========== 数字信封协议演示完成 ==========\n");
    return ret;
}